import traceback
from typing import List, Dict

from pymysql import DatabaseError

from conf.config import MYSQL_CONFIG, logger
from db.mysql.mysql_db import MysqlClient


class TextChunkDAO(object):
    """ 文本块数据访问对象 """

    def __init__(self, *args, **kwargs):
        super(TextChunkDAO, self).__init__(*args, **kwargs)

    @staticmethod
    def get_mysql_conn():
        """ 获取mysql连接 """
        mysql_conn = MysqlClient(MYSQL_CONFIG)
        return mysql_conn

    def get_text_chunks(self, chunk_ids: List[int]) -> List[Dict]:
        """
        获取文本块列表
        :param chunk_ids: 文本块的id列表
        :return:
        """
        mysql_conn = self.get_mysql_conn()

        text_chunks = list()
        try:
            sql = ('SELECT id, content, content_vector FROM text_chunk WHERE id in %s '
                   'ORDER BY id;')
            args = (chunk_ids, )
            count, text_chunks = mysql_conn.fetchall(sql, args=args)
        finally:
            if "mysql_conn" in dir():  # 判断连接是否成功创建，创建了才能执行close()
                mysql_conn.close()

        return text_chunks

    def insert_text_chunks(self, video_id: int, content: str):
        """
        插入文本块
        :param video_id: 视频id
        :param content: 内容
        :return:
        """
        mysql_conn = self.get_mysql_conn()
        try:
            sql = ('INSERT INTO `text_chunk`(video_id, content) '
                   'VALUES(%s, %s);')
            args = (video_id, content, )
            mysql_conn.execute(sql, args)

            # 提交事物
            mysql_conn.commit()
        except DatabaseError:
            mysql_conn.rollback()
            error_str = traceback.format_exc()
            logger.error(error_str)
        finally:
            if "mysql_conn" in dir():  # 判断连接是否成功创建，创建了才能执行close()
                mysql_conn.close()

    def update_content_vector(self, chunk_id: int, content_vector: str):
        """
        更新内容对应的向量值
        :param chunk_id: 主键，即id
        :param content_vector: 内容对应的向量值
        :return:
        """
        mysql_conn = self.get_mysql_conn()
        try:
            sql = ('UPDATE `text_chunk` SET content_vector = %s '
                   'WHERE id = %s;')
            args = (content_vector, chunk_id,)
            mysql_conn.execute(sql, args)

            # 提交事物
            mysql_conn.commit()
        except DatabaseError:
            mysql_conn.rollback()
            error_str = traceback.format_exc()
            logger.error(error_str)
        finally:
            if "mysql_conn" in dir():  # 判断连接是否成功创建，创建了才能执行close()
                mysql_conn.close()


def main():
    text_chunk_dao = TextChunkDAO()
    text_chunks = text_chunk_dao.get_text_chunks(chunk_ids=[1, 3])
    print(text_chunks)


if __name__ == '__main__':
    main()
